# 时间：2025-07-25
# 作者：WMSmile
# 描述：API模块，用于与ragflow服务端进行交互
import requests
from utils.logger import logger
import os


class RagflowAPI:
    def __init__(self):
        self.ragflow_api_url = os.getenv("RAGFLOW_API_URL")
        self.ragflow_kb_id = os.getenv("RAGFLOW_KB_ID")
        self.ragflow_kb_name = os.getenv("RAGFLOW_KB_NAME")
        self.ragflow_api_key = os.getenv("RAGFLOW_API_KEY")
        logger.debug(f"Ragflow API URL: {self.ragflow_api_url}")
        logger.debug(f"Ragflow KB ID: {self.ragflow_kb_id}")
        logger.debug(f"Ragflow KB Name: {self.ragflow_kb_name}")
        logger.debug(f"Ragflow API Key: {self.ragflow_api_key}")

    def get_headers(self,type=1) -> dict:
        """
        Get the headers for the Ragflow API.
        """
        if type == 2:
            return {
                "Authorization": f"Bearer {self.ragflow_api_key}",
                # "Content-Type": "multipart/form-data"
            }
        return {
            "Authorization": f"Bearer {self.ragflow_api_key}",
            "Content-Type": "application/json"
        }

    def is_succeed(self,response) -> bool:
        """
        检查API响应是否成功。
        
        根据API响应中的 'code' 字段判断请求是否成功。
        在该API的响应规范中，'code' 字段值为 0 表示请求成功。

        Args:
            response (dict): 从API接收到的响应数据，通常为字典类型。

        Returns:
            bool: 如果 'code' 字段的值为 0，则返回 True；否则返回 False。
        """
        return response.get("code") == 0

    def check_api_url(self) -> tuple[bool, str]:
        """
        检测配置的 configs.API_URL 是否可以访问
        版本：v0.19.0 full
        错误代码提示
        Code	Message	Description
        400	Bad Request	Invalid request parameters
        401	Unauthorized	Unauthorized access
        403	Forbidden	Access denied
        404	Not Found	Resource not found
        500	Internal Server Error	Server internal error
        1001	Invalid Chunk ID	Invalid Chunk ID
        1002	Chunk Update Failed	Chunk update failed
        """
        print(f"【检查API】{self.ragflow_api_url}")
        reqUrl = f"{self.ragflow_api_url}/datasets/{self.ragflow_kb_id}/documents?page=1&page_size=1"
        try:
            logger.info(f"【检查API】url = {reqUrl}")
            logger.info(f"【检查API】headers = {self.get_headers()}")
            r = requests.get(reqUrl, headers=self.get_headers())
        except Exception as e:
            return False, f"请求失败，请检查API相关配置后重试，请求异常：{e}"
        
        if r.status_code != 200:
            return False, f"请求失败，请检查API相关配置后重试，请求状态码：{r.status_code}"
        
        response = r.json()
        if self.is_succeed(response):
            logger.info(f"ragflow version：{response.get('data')}")
            return True, f"API地址配置正确"
        
        code = response.get("code")
        message = response.get("message")
        
        if code == 401 or code == 403:
            return False, "认证失败/未授权，请检查 AUTHORIZATION 配置（授权Token）"
        elif code == 100 or '404' in message:
            return False, "API地址配置错误，请检查 API_URL 配置（API地址）"
        else:
            return False, "请求失败，请检查API相关配置后重试"
    
    def upload_kb_file(self, files:list[str]) -> tuple[bool, str]:
        """
        上传知识库文件
        """
        url = f"{self.ragflow_api_url}/datasets/{self.ragflow_kb_id}/documents"
        try:
            # logger.info(f"【上传文件】{url}")
            # logger.info(f"【上传文件】{self.get_headers(2)}")
            upload_files = [("file", open(file, "rb")) for file in files]
            # upload_files = [("file", (f, open(f, "rb"), "application/octet-stream")) for f in files]
            # upload_files = [("file", (os.path.basename(f), open(f, "rb"), "application/octet-stream")) for f in files]
            # logger.info(f"【上传文件】{upload_files}")
            response = requests.post(
                url=url,
                headers=self.get_headers(2),
                files=upload_files,
            )
            # logger.info(f"【上传文件】{response.request.headers}")
            # logger.info(f"【上传文件】{response.request.body}")
            logger.info(f"【上传文件】{response.status_code}")
            logger.info(f"【上传文件】{response.text}")
            #关闭file
            for f in upload_files:
                f[1].close()
            return (True,response.json())
        except Exception as e:
            logger.error(f"【上传文件】{e}")
            return (False,f"【上传文件】{e}")

        


    